import socket
import pygame
import time
from enum import Enum

class ShipState(Enum):
    NAVIGATION = 1    # 导航状态
    REMOTE_CTRL = 2   # 遥控状态
    STOP = 3

# 蓝牙连接模式
AXIS_X = 0  # 左摇杆X轴
AXIS_Y = 1  # 左摇杆Y轴
AXIS_R_X = 3  # 右摇杆X轴
AXIS_R_Y = 4  # 右摇杆Y轴

BUTTON_A = 0  # A键
BUTTON_B = 1  # B键
BUTTON_X = 3  # X键
BUTTON_Y = 4  # Y键
LB_KEY = 6  # 左肩键
RB_KEY = 7  # 右肩键

PWM_DUTY_STOP = 7.5  # 停止时的PWM占空比

TCP_REMOTE_IP = '192.168.191.45'
TCP_REMOTE_PORT = 11889

# 配置
useJoystick = False
udp_socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
target_address = (TCP_REMOTE_IP, TCP_REMOTE_PORT)
ctrl_mode = 1  # 1：停止船只 2：启动船只
joystick = None
# 初始化pygame和手柄
def init_joystick():
    pygame.init()
    pygame.joystick.quit()  # 确保手柄模块被正确初始化
    pygame.joystick.init()
    if pygame.joystick.get_count() == 0:
        print("未检测到手柄，请连接手柄后重试。")
        return None
    print(f"检测到 {pygame.joystick.get_count()} 个手柄。")
    joystick = pygame.joystick.Joystick(0)
    joystick.init()
    return joystick

def check_joystick():
    if not pygame.joystick.get_init():
        pygame.joystick.init()
    if pygame.joystick.get_count() == 0:
        return False
    return True
    
# 遥控船状态类
class BoatState:
    def __init__(self):
        self.pwm_neutral = 7.5
        self.pwm_min = 7.00
        self.pwm_max = 8.10
        self.base_speed = 0.0
        self.turn_gain = 0.5
        self.max_throttle = 1.0
        self.boost_multiplier = 1.5

boat = BoatState()

def map_speed_to_pwm(speed):
    if speed >= 0:
        pwm = boat.pwm_neutral + speed * (boat.pwm_max - boat.pwm_neutral)
    else:
        pwm = boat.pwm_neutral + speed * (boat.pwm_neutral - boat.pwm_min)
    return max(boat.pwm_min, min(pwm, boat.pwm_max))

def update_esc_pwm(left_speed, right_speed):
    left_pwm = round(map_speed_to_pwm(left_speed),1)
    right_pwm = round(map_speed_to_pwm(right_speed),1) 
    return left_pwm, right_pwm

def handle_joystick_input(joystick):
    y_axis = -joystick.get_axis(AXIS_Y)
    x_axis = joystick.get_axis(AXIS_R_X)
    a_button = joystick.get_button(BUTTON_A)

    boat.base_speed = y_axis
    current_max_throttle = boat.max_throttle * (boat.boost_multiplier if a_button else 1.0)
    if a_button:
        boat.pwm_max = boat.pwm_max+0.05 if boat.pwm_max<=9.5 else 9.5
    else:
        boat.pwm_max = 8.10
    boat.base_speed = max(-current_max_throttle, min(boat.base_speed, current_max_throttle))

    turn_effect = x_axis * boat.turn_gain
    left_speed = boat.base_speed  + turn_effect
    right_speed = boat.base_speed - turn_effect

    left_speed = max(-current_max_throttle, min(left_speed, current_max_throttle))
    right_speed = max(-current_max_throttle, min(right_speed, current_max_throttle))

    left_pwm, right_pwm = update_esc_pwm(left_speed, right_speed)
    print(f"L_PWM: {left_pwm:.1f}% | R_PWM: {right_pwm:.1f}% | Speed: {boat.base_speed:.2f}")
    return left_pwm, right_pwm

def map_pwm(pwm0_min, pwm0_max, pwm_min, pwm_max, pwm_cur):
    pass
    

def send_control_message(ctrl_mode, left_pwm, right_pwm):
    msg = f'{ctrl_mode},{int(left_pwm*10)},{int(right_pwm*10)}'
    udp_socket.sendto(msg.encode('utf-8'), target_address)

# 主循环
try:
    while True:
        if joystick is None:
            print(time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()), "等待手柄连接...")
            joystick = init_joystick()
            time.sleep(1)
            continue
        if not check_joystick():
            print("手柄未连接或初始化失败，请检查连接。")
            joystick = None
            useJoystick = False
            time.sleep(1)
            continue
        pygame.event.pump()
        if useJoystick:
            left_pwm, right_pwm = handle_joystick_input(joystick)
            send_control_message(ctrl_mode, left_pwm, right_pwm)

        if joystick.get_button(RB_KEY):  # 停止船只
            ctrl_mode = ShipState.STOP.value
            left_pwm, right_pwm = handle_joystick_input(joystick)
            send_control_message(ctrl_mode, PWM_DUTY_STOP, PWM_DUTY_STOP)
            useJoystick = False

        if joystick.get_button(LB_KEY):  # 启动船只
            ctrl_mode = ShipState.REMOTE_CTRL.value
            useJoystick = True

        time.sleep(0.1)
except KeyboardInterrupt:
    print("程序已终止。")
finally:
    pygame.quit()